#!/usr/bin/env python
"""

    run_neighbor.py

"""

################################################################################
#
#   run_neighbor
# 
#       from phylip package
# 
# 
#   Copyright (C) 2009 Leo Goodstadt
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; version 2
#   of the License
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#################################################################################

import sys, os
from collections import defaultdict
import json
#import popen2
import subprocess
import tempfile
import string
import shutil
import glob
    


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#_________________________________________________________________________________________
# 
#   run_neighbor
# 
#_________________________________________________________________________________________
def _write_matrix_file (filename, matrix_names, matrix):
    """
    returns file name of written sequence
    """
    matrix_file = open(filename, "w")
    matrix_file.write("%d\n" % len(matrix_names))
    cnt_items = len(matrix_names)
    for ii in range(cnt_items):
        matrix_file.write(matrix_names[ii] + "\t")
        matrix_file.write("\t".join (map(str,  matrix[ii])) + "\n")
        
    
    
    
    
def run_neighbor (matrix_names, matrix, neighbor_bin = "neighbor"):
    """
    Calculates a neighbor joined tree from a matrix
    
    returns a NH format tree in a string

    """
    try:
        if len(matrix) != len(matrix[0]):
            raise Exception("Expected square matrix (%d x %d)" % 
                                                        (len(matrix), len(matrix[0])) )
        if len(matrix) != len(matrix_names):
            raise Exception("Expected matrix size (%d) = matrix_names size (%d)" % 
                                                        (len(matrix),  len(matrix_names)))

        temp_matrix_names = ["XY%08d" % n for n in range(len(matrix))]
        temp_dir = tempfile.mkdtemp()
        _write_matrix_file(os.path.join(temp_dir, "infile"), temp_matrix_names, matrix)
        
        #sys.stderr.write(open(os.path.join(temp_dir, "infile")).read())
        
        process = subprocess.Popen(neighbor_bin,
                                   stdin=subprocess.PIPE,
                                   stderr=subprocess.PIPE,
                                   stdout=subprocess.PIPE,
                                   cwd=temp_dir,
                                   shell=True
                                   )    
        (stdout_str, stderr_str) = process.communicate("Y\n")

        tree_output_file_name = os.path.join(temp_dir, "outtree")

        #
        #   failed: throw exception
        # 
        if (process.returncode != 0 or 
            not os.path.exists(tree_output_file_name)):
            err_str = ("neighbor failed to run, returning %s\n%s\n%s\n" % (process.returncode, 
                                                                            stdout_str,
                                                                            stderr_str))
            raise Exception(err_str)


        tree_nh_str = "".join(open(tree_output_file_name).readlines()).rstrip()
        
        shutil.rmtree(temp_dir)

        for temp_name, orig_name in zip(temp_matrix_names,  matrix_names):
            tree_nh_str = tree_nh_str.replace(temp_name, orig_name)
        return tree_nh_str
    except:
        shutil.rmtree(temp_dir)
        raise

   
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
import unittest
class Test_run_neighbor_size(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def test_function(self):
        """
            test 
        """
        import re
        from collections import defaultdict
        
        for pairs_file in glob.glob("test/test_run_gencompress_size/*.pairs"):
            values = defaultdict(dict)
            for line in open(pairs_file):
                fields = line.rstrip().split("\t")
                values[fields[0]][fields[1]] = fields[2]
                values[fields[1]][fields[0]] = fields[2]
                
                
            matrix_size = len(values.keys())
            matrix = [[0] * matrix_size for m in range(matrix_size)]
            matrix_names_values = dict()
            matrix_names = []
            for ii, k1 in enumerate(values.keys()):
                matrix_names.append(k1)
                matrix_names_values[k1] = ii
            for ii, k1 in enumerate(values.keys()):
                for k2, v in values[k1].iteritems():
                    jj = matrix_names_values[k2]
                    matrix[ii][jj] = v
                    matrix[jj][ii] = v
                    
                    
            tree_str = run_neighbor (matrix_names, matrix, )
            nj_tree_file = pairs_file.replace(".pairs", ".nj.tree")
            open(nj_tree_file,  "w").write(tree_str)



#
#   debug code not run if called as a module
#     
if __name__ == '__main__':
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()



    
